import os
import base64
from mistralai import Mistral

# --- 配置 ---
API_KEY = "Kz2yZgyEQxkreQRagET0rtz15l8smbRT"
IMAGE_PATH = "可达鸭.jpg"

# 支持视觉的模型列表
VISION_MODELS = [
    "pixtral-large-latest",  # Pixtral 视觉模型
    "mistral-small-2506",  # Mistral Small 支持视觉
    "mistral-medium-3",  # 尝试 Medium 3
    "mistral-large-latest",
    
]

def encode_image(image_path):
    """将图片编码为 base64"""
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    except FileNotFoundError:
        print(f"错误: 找不到文件 {image_path}")
        return None
    except Exception as e:
        print(f"错误: {e}")
        return None

def test_image_recognition(model_name, base64_image):
    """测试指定模型的图像识别能力"""
    print(f"\n{'='*60}")
    print(f"测试模型: {model_name}")
    print(f"{'='*60}")
    
    try:
        client = Mistral(api_key=API_KEY)
        
        # 构建消息 - 使用 base64 编码的图片
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "请描述这张图片中的内容,这是什么?"
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image}"
                    }
                ]
            }
        ]
        
        # 发送请求
        response = client.chat.complete(
            model=model_name,
            messages=messages,
            max_tokens=500
        )
        
        # 打印结果
        print(f"✓ 成功! 模型 {model_name} 支持图像识别")
        print(f"回复: {response.choices[0].message.content}")
        return True
        
    except Exception as e:
        error_msg = str(e)
        print(f"✗ 失败! 模型 {model_name} 不支持图像或出现错误")
        print(f"错误信息: {error_msg}")
        
        # 判断是否是不支持图像的错误
        if "image" in error_msg.lower() or "vision" in error_msg.lower() or "multimodal" in error_msg.lower():
            print("→ 该模型可能不支持图像输入")
        
        return False

def test_agent_stream_with_image():
    """测试 Agent Stream 模式是否支持图像"""
    print(f"\n{'='*60}")
    print(f"测试模型: mistral-agent-stream (Agent ID)")
    print(f"{'='*60}")
    
    try:
        client = Mistral(api_key=API_KEY)
        base64_image = encode_image(IMAGE_PATH)
        
        if not base64_image:
            return False
        
        # Agent 模式使用 inputs 参数
        inputs = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "请描述这张图片中的内容"
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image}"
                    }
                ]
            }
        ]
        
        response_stream = client.beta.conversations.start_stream(
            inputs=inputs,
            agent_id="ag_019a3501501b768a901bef17061c2a20",
        )
        
        print("✓ Agent Stream 模式支持图像!")
        print("回复: ", end="", flush=True)
        
        for chunk in response_stream:
            if chunk.event == 'message.output.delta':
                print(chunk.data.content, end="", flush=True)
        
        print()
        return True
        
    except Exception as e:
        print(f"✗ Agent Stream 模式不支持图像或出现错误")
        print(f"错误信息: {e}")
        return False

# --- 主程序 ---
if __name__ == "__main__":
    print("="*60)
    print("Mistral AI 图像识别能力测试")
    print("="*60)
    
    # 检查图片文件是否存在
    if not os.path.exists(IMAGE_PATH):
        print(f"错误: 找不到图片文件 {IMAGE_PATH}")
        print("请确保 '可达鸭.jpg' 文件在当前目录下")
        exit(1)
    
    # 编码图片
    print(f"\n正在读取图片: {IMAGE_PATH}")
    base64_image = encode_image(IMAGE_PATH)
    
    if not base64_image:
        print("图片编码失败,退出测试")
        exit(1)
    
    print(f"✓ 图片编码成功 (大小: {len(base64_image)} 字符)")
    
    # 测试各个模型
    results = {}
    for model in VISION_MODELS:
        results[model] = test_image_recognition(model, base64_image)
    
    # 测试 Agent Stream 模式
    results["mistral-agent-stream"] = test_agent_stream_with_image()
    
    # 汇总结果
    print(f"\n{'='*60}")
    print("测试结果汇总")
    print(f"{'='*60}")
    
    for model, success in results.items():
        status = "✓ 支持图像" if success else "✗ 不支持图像"
        print(f"{model:30s} : {status}")
    
    print(f"\n{'='*60}")
    print("测试完成!")
    print(f"{'='*60}")